#include "graph.h"
#include <time.h>
#include <limits.h>

#ifdef _OPENMP
#include <omp.h>
#endif

// ========================================================================
// 优化算法实现
// ========================================================================

// 复制调度方案
Schedule* copy_schedule(const Schedule* original) {
    if (!original) return NULL;
    
    Schedule* copy = create_schedule(original->length);
    if (!copy) return NULL;
    
    memcpy(copy->schedule, original->schedule, original->length * sizeof(int));
    memcpy(copy->vstay_history, original->vstay_history, original->length * sizeof(int));
    copy->max_vstay = original->max_vstay;
    
    return copy;
}

// 计算节点的合法移动范围（完全匹配MATLAB逻辑）
typedef struct {
    int* positions;
    int count;
} ValidPositions;

ValidPositions* calculate_valid_move_range(const int* schedule, int length, 
                                         const Graph* g, int node_idx_to_move) {
    int node_to_move = schedule[node_idx_to_move];
    
    // === 对应MATLAB: ancestors = find(is_ancestor(:, node_to_move)) ===
    int* ancestors = malloc(length * sizeof(int));
    int ancestor_count = 0;
    for (int i = 0; i < g->num_nodes; i++) {
        if (g->is_ancestor[i][node_to_move]) {
            ancestors[ancestor_count++] = i;
        }
    }
    
    // === 对应MATLAB: descendants = find(is_ancestor(node_to_move, :)) ===
    int* descendants = malloc(length * sizeof(int));
    int descendant_count = 0;
    for (int i = 0; i < g->num_nodes; i++) {
        if (g->is_ancestor[node_to_move][i]) {
            descendants[descendant_count++] = i;
        }
    }
    
    // === 对应MATLAB: ancestor_positions = positions_in_schedule(ismember(local_schedule, ancestors)) ===
    int* ancestor_positions = NULL;
    int ancestor_pos_count = 0;
    if (ancestor_count > 0) {
        ancestor_positions = malloc(ancestor_count * sizeof(int));
        for (int i = 0; i < length; i++) {
            int node_in_schedule = schedule[i];
            // 检查是否在祖先列表中
            for (int j = 0; j < ancestor_count; j++) {
                if (node_in_schedule == ancestors[j]) {
                    ancestor_positions[ancestor_pos_count++] = i + 1; // MATLAB使用1-based索引
                    break;
                }
            }
        }
    }
    
    // === 对应MATLAB: latest_ancestor_pos = max(ancestor_positions) ===
    int latest_ancestor_pos = 0; // MATLAB默认0
    if (ancestor_positions) {
        for (int i = 0; i < ancestor_pos_count; i++) {
            if (ancestor_positions[i] > latest_ancestor_pos) {
                latest_ancestor_pos = ancestor_positions[i];
            }
        }
    }
    
    // === 对应MATLAB: descendant_positions = positions_in_schedule(ismember(local_schedule, descendants)) ===
    int* descendant_positions = NULL;
    int descendant_pos_count = 0;
    if (descendant_count > 0) {
        descendant_positions = malloc(descendant_count * sizeof(int));
        for (int i = 0; i < length; i++) {
            int node_in_schedule = schedule[i];
            // 检查是否在后代列表中
            for (int j = 0; j < descendant_count; j++) {
                if (node_in_schedule == descendants[j]) {
                    descendant_positions[descendant_pos_count++] = i + 1; // MATLAB使用1-based索引
                    break;
                }
            }
        }
    }
    
    // === 对应MATLAB: earliest_descendant_pos = min(descendant_positions) ===
    int earliest_descendant_pos = length + 1; // MATLAB默认num_nodes + 1
    if (descendant_positions) {
        for (int i = 0; i < descendant_pos_count; i++) {
            if (descendant_positions[i] < earliest_descendant_pos) {
                earliest_descendant_pos = descendant_positions[i];
            }
        }
    }
    
    // === 对应MATLAB: valid_range = (latest_ancestor_pos + 1):(earliest_descendant_pos - 1) ===
    ValidPositions* result = malloc(sizeof(ValidPositions));
    result->positions = malloc(length * sizeof(int));
    result->count = 0;
    
    int range_start = latest_ancestor_pos + 1;
    int range_end = earliest_descendant_pos - 1;
    
    for (int pos = range_start; pos <= range_end; pos++) {
        // === 对应MATLAB: valid_range(valid_range == node_idx_to_move) = [] ===
        if (pos != (node_idx_to_move + 1)) { // 转换为1-based索引比较
            result->positions[result->count++] = pos - 1; // 转换回0-based索引
        }
    }
    
    // 清理临时数组
    free(ancestors);
    free(descendants);
    if (ancestor_positions) free(ancestor_positions);
    if (descendant_positions) free(descendant_positions);
    
    return result;
}

void free_valid_positions(ValidPositions* vp) {
    if (vp) {
        free(vp->positions);
        free(vp);
    }
}

// 生成节点移动后的新调度
int* generate_moved_schedule(const int* original, int length, 
                            int node_idx, int new_pos) {
    int* new_schedule = malloc(length * sizeof(int));
    if (!new_schedule) return NULL;
    
    int node_to_move = original[node_idx];
    
    // 生成移除节点后的临时序列
    int temp_idx = 0;
    for (int i = 0; i < length; i++) {
        if (i != node_idx) {
            new_schedule[temp_idx++] = original[i];
        }
    }
    
    // 插入节点到新位置
    int adjusted_pos = new_pos;
    if (new_pos > node_idx) adjusted_pos--; // 因为移除了一个元素
    
    // 向右移动元素为新节点腾出空间
    for (int i = length - 2; i >= adjusted_pos; i--) {
        new_schedule[i + 1] = new_schedule[i];
    }
    new_schedule[adjusted_pos] = node_to_move;
    
    return new_schedule;
}

// ========================================================================
// 阶段一优化器：随机采样 + 瓶颈攻击
// ========================================================================
Schedule* optimizer_random_sampling(const Schedule* initial, const Graph* g, 
                                   const OptOptions* opts, const char* output_dir) {
    printf("Starting Phase 1 optimization (random sampling)...\n");
    
    Schedule* current = copy_schedule(initial);
    Schedule* best = copy_schedule(initial);
    int iterations_without_improvement = 0;
    
    double start_time = get_time_seconds();
    
    // === 创建详细收敛日志文件（记录每次迭代） ===
    char convergence_log_path[512];
    snprintf(convergence_log_path, sizeof(convergence_log_path), "%s/iteration_convergence.csv", output_dir);
    FILE* convergence_log = fopen(convergence_log_path, "w");
    if (convergence_log) {
        fprintf(convergence_log, "Iteration,V_stay,Phase,Time_Seconds,Improvement_Flag,Description\n");
        fprintf(convergence_log, "0,%d,Phase1,0.00,0,Initial solution\n", initial->max_vstay);
        fflush(convergence_log);
    }
    
    for (int iter = 0; iter < opts->iterations; iter++) {
        // 1. 瓶颈攻击：找到峰值内存位置
        int peak_idx = 0;
        for (int i = 1; i < current->length; i++) {
            if (current->vstay_history[i] > current->vstay_history[peak_idx]) {
                peak_idx = i;
            }
        }
        
        // 搜索空间：从开始到峰值位置
        int search_space_size = peak_idx + 1;
        if (search_space_size == 0) search_space_size = current->length;
        
        // 2. 批量生成候选解
        int batch_size = opts->batch_size;
        int* best_candidate = NULL;
        int best_candidate_score = INT_MAX;
        
        #ifdef _OPENMP
        #pragma omp parallel
        #endif
        {
            int local_best_score = INT_MAX;
            int* local_best_candidate = NULL;
            
            #ifdef _OPENMP
            #pragma omp for
            #endif
            for (int k = 0; k < batch_size; k++) {
                // === 对应MATLAB: node_idx_to_move = search_space_indices(randi(length(search_space_indices))) ===
                int node_idx_to_move = rand() % search_space_size;
                int node_to_move = current->schedule[node_idx_to_move];
                
                // === 使用新的合法移动范围计算（完全匹配MATLAB逻辑） ===
                ValidPositions* valid_pos = calculate_valid_move_range(current->schedule, 
                                                                      current->length, 
                                                                      g, 
                                                                      node_idx_to_move);
                
                if (!valid_pos || valid_pos->count == 0) {
                    if (valid_pos) free_valid_positions(valid_pos);
                    continue;
                }
                
                // === 对应MATLAB: new_pos = valid_range(randi(length(valid_range))) ===
                int new_pos = valid_pos->positions[rand() % valid_pos->count];
                
                // === 对应MATLAB的节点移动逻辑 ===
                // temp_schedule = local_schedule;
                // temp_schedule(node_idx_to_move) = [];
                // if new_pos > node_idx_to_move, new_pos = new_pos - 1; end
                // temp_schedule = [temp_schedule(1:new_pos-1), node_to_move, temp_schedule(new_pos:end)];
                
                int* temp_schedule = malloc(current->length * sizeof(int));
                if (!temp_schedule) {
                    free_valid_positions(valid_pos);
                    continue;
                }
                
                // 移除要移动的节点，生成临时序列
                int temp_idx = 0;
                for (int i = 0; i < current->length; i++) {
                    if (i != node_idx_to_move) {
                        temp_schedule[temp_idx++] = current->schedule[i];
                    }
                }
                
                // 调整插入位置（对应MATLAB的位置调整逻辑）
                int adjusted_pos = new_pos;
                if (new_pos > node_idx_to_move) {
                    adjusted_pos = new_pos - 1;
                }
                
                // 在指定位置插入节点
                for (int i = current->length - 2; i >= adjusted_pos; i--) {
                    temp_schedule[i + 1] = temp_schedule[i];
                }
                temp_schedule[adjusted_pos] = node_to_move;
                
                free_valid_positions(valid_pos);
                
                // === 对应MATLAB: is_l0_schedule_valid + 拓扑依赖检查 ===
                if (is_schedule_fully_valid(temp_schedule, current->length, g)) {
                    int score = calculate_max_vstay(temp_schedule, current->length, g, NULL);
                    
                    if (score < local_best_score) {
                        if (local_best_candidate) free(local_best_candidate);
                        local_best_candidate = temp_schedule;
                        local_best_score = score;
                        temp_schedule = NULL; // 防止被释放
                    }
                }
                
                if (temp_schedule) free(temp_schedule);
            }
            
            // 合并线程结果
            #ifdef _OPENMP
            #pragma omp critical
            #endif
            {
                if (local_best_score < best_candidate_score) {
                    if (best_candidate) free(best_candidate);
                    best_candidate = local_best_candidate;
                    best_candidate_score = local_best_score;
                    local_best_candidate = NULL; // 防止被释放
                }
                if (local_best_candidate) free(local_best_candidate);
            }
        } // end parallel
        
        // 3. 更新最优解
        if (best_candidate_score < best->max_vstay) {
            free(current->schedule);
            free(current->vstay_history);
            
            current->schedule = best_candidate;
            // 重新分配vstay_history数组
            current->vstay_history = malloc(current->length * sizeof(int));
            current->max_vstay = calculate_max_vstay(best_candidate, current->length, 
                                                   g, current->vstay_history);
            
            if (current->max_vstay < best->max_vstay) {
                free_schedule(best);
                best = copy_schedule(current);
            }
            
            iterations_without_improvement = 0;
            printf("  Phase 1 > Iteration %d/%d: Found better solution! New score: %d\n", 
                   iter + 1, opts->iterations, current->max_vstay);
            best_candidate = NULL; // 防止被释放
        } else {
            iterations_without_improvement++;
            if (iter % 100 == 0) {
                printf("  Phase 1 > Iteration %d/%d: Current best: %d, no improvement for: %d\n",
                       iter + 1, opts->iterations, best->max_vstay, 
                       iterations_without_improvement);
            }
        }
        
        if (best_candidate) free(best_candidate);
        
        // === 记录每次迭代的收敛数据 ===
        if (convergence_log) {
            double elapsed_time = get_time_seconds() - start_time;
            int improvement_flag = (best_candidate_score < best->max_vstay && best_candidate_score != INT_MAX) ? 1 : 0;
            fprintf(convergence_log, "%d,%d,Phase1,%.2f,%d,Iteration %d\n", 
                    iter + 1, best->max_vstay, elapsed_time, improvement_flag, iter + 1);
            fflush(convergence_log);
        }
        
        // 4. 提前终止检查
        if (iterations_without_improvement >= opts->patience) {
            printf("  阶段一 > 连续 %d 次迭代未改进，提前终止\n", opts->patience);
            break;
        }
    }
    
    double elapsed = get_time_seconds() - start_time;
    printf("Phase 1 completed! Time: %.2f seconds\n", elapsed);
    
    // === 不在这里关闭日志文件，Phase 2还需要使用 ===
    
    free_schedule(current);
    return best;
}

// 生成节点交换后的新调度（对应MATLAB的交换逻辑）
int* generate_swapped_schedule(const int* original, int length, int node_A_id, int node_B_id) {
    int* new_schedule = malloc(length * sizeof(int));
    if (!new_schedule) return NULL;
    
    // 复制原调度
    memcpy(new_schedule, original, length * sizeof(int));
    
    // 找到两个节点在调度中的位置
    int pos_A = -1, pos_B = -1;
    for (int i = 0; i < length; i++) {
        if (original[i] == node_A_id) pos_A = i;
        if (original[i] == node_B_id) pos_B = i;
    }
    
    // 如果找到了两个位置，执行交换
    if (pos_A >= 0 && pos_B >= 0) {
        new_schedule[pos_A] = node_B_id;
        new_schedule[pos_B] = node_A_id;
    }
    
    return new_schedule;
}

// 检查两个节点是否可以交换（无依赖关系）
bool can_nodes_swap(const Graph* g, int node_A_id, int node_B_id) {
    // 两个节点可以交换当且仅当它们之间没有依赖关系
    return !g->is_ancestor[node_A_id][node_B_id] && !g->is_ancestor[node_B_id][node_A_id];
}

// 为焦点节点生成所有合法的交换伙伴
typedef struct {
    int node_A;
    int node_B;
} SwapPair;

typedef struct {
    SwapPair* pairs;
    int count;
    int capacity;
} SwapCandidates;

SwapCandidates* generate_swap_candidates(const int* schedule, int length, 
                                        const int* focus_nodes, int focus_count,
                                        const Graph* g) {
    SwapCandidates* candidates = malloc(sizeof(SwapCandidates));
    candidates->capacity = focus_count * length;
    candidates->pairs = malloc(candidates->capacity * sizeof(SwapPair));
    candidates->count = 0;
    
    for (int i = 0; i < focus_count; i++) {
        int focus_node = focus_nodes[i];
        
        // 为每个焦点节点寻找交换伙伴
        for (int j = 0; j < length; j++) {
            int partner_node = schedule[j];
            if (focus_node == partner_node) continue;
            
            // 检查是否可以交换（无依赖关系）
            if (can_nodes_swap(g, focus_node, partner_node)) {
                if (candidates->count < candidates->capacity) {
                    candidates->pairs[candidates->count].node_A = focus_node;
                    candidates->pairs[candidates->count].node_B = partner_node;
                    candidates->count++;
                }
            }
        }
    }
    
    return candidates;
}

void free_swap_candidates(SwapCandidates* candidates) {
    if (candidates) {
        free(candidates->pairs);
        free(candidates);
    }
}

// ========================================================================
// 阶段二优化器：节点交换策略（完全匹配MATLAB逻辑）
// ========================================================================
Schedule* optimizer_systematic_search(const Schedule* initial, const Graph* g, 
                                     const OptOptions* opts, const char* output_dir) {
    printf("Starting Phase 2 optimization (systematic search - node swapping)...\n");
    
    Schedule* current = copy_schedule(initial);
    Schedule* best = copy_schedule(initial);
    
    // 获取CPU核心数
    int num_cores = 1;
    #ifdef _OPENMP
    num_cores = omp_get_max_threads();
    #endif
    
    double start_time = get_time_seconds();
    
    // === 打开收敛日志文件（追加模式） ===
    char convergence_log_path[512];
    snprintf(convergence_log_path, sizeof(convergence_log_path), "%s/iteration_convergence.csv", output_dir);
    FILE* convergence_log = fopen(convergence_log_path, "a");
    
    // 从文件中读取最后的迭代次数
    int last_iteration = 0;
    FILE* read_log = fopen(convergence_log_path, "r");
    if (read_log) {
        char line[256];
        while (fgets(line, sizeof(line), read_log)) {
            int iter;
            if (sscanf(line, "%d,", &iter) == 1) {
                if (iter > last_iteration) last_iteration = iter;
            }
        }
        fclose(read_log);
    }
    
    for (int pass_num = 0; pass_num < opts->max_passes; pass_num++) {
        // 时间限制检查（对应MATLAB的时间检查）
        double elapsed = get_time_seconds() - start_time;
        if (elapsed > opts->time_limit_minutes * 60.0) {
            printf("  Phase 2 > Time limit of %d minutes reached, terminating optimization.\n", 
                   opts->time_limit_minutes);
            break;
        }
        
        printf("  Phase 2 > Starting round %d / %d (current best: %d)\n", 
               pass_num + 1, opts->max_passes, current->max_vstay);
        
        bool improvement_found_in_pass = false;
        
        // === 稳健的候选生成机制（对应MATLAB的候选生成循环） ===
        SwapCandidates* candidate_pool = NULL;
        int attempts = 0;
        
        // 循环直到成功生成候选，或尝试次数过多
        while ((!candidate_pool || candidate_pool->count == 0) && attempts < 10) {
            attempts++;
            
            if (candidate_pool) {
                free_swap_candidates(candidate_pool);
                candidate_pool = NULL;
            }
            
            // 1. 瓶颈攻击: 从峰值前序列选择"焦点"节点
            int peak_idx = 0;
            for (int i = 1; i < current->length; i++) {
                if (current->vstay_history[i] > current->vstay_history[peak_idx]) {
                    peak_idx = i;
                }
            }
            
            int search_space_size = peak_idx + 1;
            if (search_space_size == 0) search_space_size = current->length;
            
            // 选择焦点节点数量（对应MATLAB的num_focus_nodes计算）
            int num_focus_nodes = (num_cores < search_space_size) ? num_cores : search_space_size;
            if (opts->num_threads > 0) {
                num_focus_nodes = (opts->num_threads < search_space_size) ? 
                                 opts->num_threads : search_space_size;
            }
            
            // 随机选择焦点节点
            int* focus_nodes_indices = malloc(search_space_size * sizeof(int));
            for (int i = 0; i < search_space_size; i++) {
                focus_nodes_indices[i] = i;
            }
            
            // Fisher-Yates洗牌算法
            for (int i = search_space_size - 1; i > 0; i--) {
                int j = rand() % (i + 1);
                int temp = focus_nodes_indices[i];
                focus_nodes_indices[i] = focus_nodes_indices[j];
                focus_nodes_indices[j] = temp;
            }
            
            // 提取焦点节点ID
            int* focus_nodes = malloc(num_focus_nodes * sizeof(int));
            for (int i = 0; i < num_focus_nodes; i++) {
                focus_nodes[i] = current->schedule[focus_nodes_indices[i]];
            }
            
            // 2. 生成所有合法交换对
            candidate_pool = generate_swap_candidates(current->schedule, current->length,
                                                     focus_nodes, num_focus_nodes, g);
            
            free(focus_nodes_indices);
            free(focus_nodes);
        }
        
        if (!candidate_pool || candidate_pool->count == 0) {
            printf("    > Warning: No swappable node pairs found after %d attempts. Skipping this round.\n", attempts);
            if (candidate_pool) free_swap_candidates(candidate_pool);
            continue;
        }
        
        printf("    > Successfully generated %d swap candidate pairs, starting evaluation...\n", candidate_pool->count);
        
        // 随机打乱候选交换对（对应MATLAB的shuffled_candidates）
        for (int i = candidate_pool->count - 1; i > 0; i--) {
            int j = rand() % (i + 1);
            SwapPair temp = candidate_pool->pairs[i];
            candidate_pool->pairs[i] = candidate_pool->pairs[j];
            candidate_pool->pairs[j] = temp;
        }
        
        // === 分批并行评估与抢先更新 ===
        int batch_size = (opts->batch_size > 0) ? opts->batch_size : (10 * num_cores);
        
        for (int batch_start = 0; batch_start < candidate_pool->count; batch_start += batch_size) {
            int batch_end = (batch_start + batch_size < candidate_pool->count) ? 
                           batch_start + batch_size : candidate_pool->count;
            int batch_count = batch_end - batch_start;
            
            int* batch_scores = malloc(batch_count * sizeof(int));
            int** batch_schedules = malloc(batch_count * sizeof(int*));
            
            // 初始化
            for (int i = 0; i < batch_count; i++) {
                batch_scores[i] = INT_MAX;
                batch_schedules[i] = NULL;
            }
            
            // 并行评估当前批次
            #ifdef _OPENMP
            #pragma omp parallel for
            #endif
            for (int k = 0; k < batch_count; k++) {
                SwapPair pair = candidate_pool->pairs[batch_start + k];
                
                // 生成交换后的调度
                int* temp_schedule = generate_swapped_schedule(current->schedule, 
                                                             current->length,
                                                             pair.node_A, 
                                                             pair.node_B);
                if (!temp_schedule) continue;
                
                // 检查完整约束 (L0约束 + 拓扑依赖)
                if (is_schedule_fully_valid(temp_schedule, current->length, g)) {
                    batch_schedules[k] = temp_schedule;
                    batch_scores[k] = calculate_max_vstay(temp_schedule, current->length, g, NULL);
                } else {
                    free(temp_schedule);
                }
            }
            
            // 找到批次中的最优解
            int best_batch_score = INT_MAX;
            int best_batch_idx = -1;
            
            for (int k = 0; k < batch_count; k++) {
                if (batch_scores[k] < best_batch_score) {
                    best_batch_score = batch_scores[k];
                    best_batch_idx = k;
                }
            }
            
            // === 抢先更新逻辑（对应MATLAB的立即更新并重启） ===
            if (best_batch_score < current->max_vstay) {
                printf("    > Found better swap, score improved from %d -> %d\n", 
                       current->max_vstay, best_batch_score);
                
                // 更新当前解
                free(current->schedule);
                free(current->vstay_history);
                
                current->schedule = batch_schedules[best_batch_idx];
                // 重新分配 vstay_history 缓冲区，避免使用已释放的指针
                current->vstay_history = (int*)malloc((size_t)current->length * sizeof(int));
                if (!current->vstay_history) {
                    fprintf(stderr, "Error: failed to allocate vstay_history in Phase 2 update.\n");
                    // 回退：计算时传入 NULL（函数内需做判空），并继续流程
                }
                current->max_vstay = calculate_max_vstay(current->schedule, current->length, 
                                                       g, current->vstay_history);
                batch_schedules[best_batch_idx] = NULL; // 防止被释放
                
                // 更新全局最优解
                if (current->max_vstay < best->max_vstay) {
                    free_schedule(best);
                    best = copy_schedule(current);
                }
                
                
                improvement_found_in_pass = true;
                
                // 清理其他候选解
                for (int i = 0; i < batch_count; i++) {
                    if (batch_schedules[i]) free(batch_schedules[i]);
                }
                free(batch_scores);
                free(batch_schedules);
                
                break; // 跳出批次循环，进入下一轮pass
            }
            
            // 清理当前批次
            for (int i = 0; i < batch_count; i++) {
                if (batch_schedules[i]) free(batch_schedules[i]);
            }
            free(batch_scores);
            free(batch_schedules);
        } // end batch loop
        
        free_swap_candidates(candidate_pool);
        
        // === 记录每轮的收敛数据 ===
        if (convergence_log) {
            double elapsed_time = get_time_seconds() - start_time;
            int current_iteration = last_iteration + pass_num + 1;
            int improvement_flag = improvement_found_in_pass ? 1 : 0;
            fprintf(convergence_log, "%d,%d,Phase2,%.2f,%d,Round %d\n", 
                    current_iteration, best->max_vstay, elapsed_time, improvement_flag, pass_num + 1);
            fflush(convergence_log);
        }
        
        // 如果完整一轮都没有任何改进，记录但继续其他轮次
        if (!improvement_found_in_pass) {
            printf("  Phase 2 > Round %d: No improvement found, continuing with next rounds.\n", pass_num + 1);
        } else {
            printf("  Phase 2 > Round %d: Found improvements, score now: %d\n", pass_num + 1, current->max_vstay);
        }
    } // end pass loop
    
    double elapsed = get_time_seconds() - start_time;
    printf("Phase 2 completed! Time: %.2f seconds\n", elapsed);
    
    // === 关闭收敛日志文件 ===
    if (convergence_log) {
        fclose(convergence_log);
    }
    
    free_schedule(current);
    return best;
}
